import argparse
from collections import defaultdict

def read_bed(input_file):
    sites=set()
    with open(input_file, 'r') as rf:
        for line in rf:
            words = line.strip().split("\t")
            sites.add((words[0],words[1],words[5]))
    return sites

def change_dataset(input_file,output_file,singleton):
    train_fh = open(output_file, "w", buffering=512)
    singleton_num=0
    non_singleton_num=0
    with open(input_file, "r") as f:
        for line in f:
            words = line.strip().split()
            key=(words[0],words[1],words[2])
            if key in singleton:
                wline=line.rstrip('\n')[:-1]+str(1)+'\n'
                singleton_num+=1
            else:
                wline=line.rstrip('\n')[:-1]+str(0)+'\n'
                non_singleton_num+=1
            train_fh.write(wline)
    print(f'input_file has singleton {singleton_num} and non_singleton {non_singleton_num}')


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--singleton",
        type=str,
        default="/public/home/hpc234701005/data/R10.4/20221109_1654_5D_PAG68757_39c39833/train1/singleton.bisulfite.bed",
        help="singleton address.",
    )
    parser.add_argument(
        "--input",
        type=str,
        default="/public/home/hpc234701005/data/R10.4/20221109_1654_5D_PAG68757_39c39833/train1/fast5s.CG.features.21mer.tsv",
        help="trainset input address.",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="/public/home/hpc234701005/data/R10.4/20221109_1654_5D_PAG68757_39c39833/train1/fast5s.CG.features.21mer.tag.tsv",
        help="trainset output address.",
    )
    
    return parser.parse_args()

if __name__ == '__main__':
    args=parse_args()
    sites=read_bed(args.singleton)
    change_dataset(args.input,args.output,sites)
